import cv2
import numpy as np
import tensorflow as tf
from tensorflow.python.client import device_lib
import Layers


flag_log_dir = './log'
NUM_CLASSES = 5
tfob_sess = None
feed = None
outputs = None


# 数据组织
def construct_data():
	tfph_data = tf.placeholder(dtype = tf.uint8, shape = [32, 32], name = 'ph_data')
	input_data = tfph_data
	input_data = tf.expand_dims(input_data, -1)
	input_data = tf.image.convert_image_dtype(input_data, dtype = tf.float32)
	input_data = tf.subtract(input_data, 0.5)
	input_data = tf.multiply(input_data, 2.0)
	input_data = tf.expand_dims(input_data, 0)

	return tfph_data, input_data


# 网络结构
def network(data, tfv_train_phase = None, name = None):
	if name is None:
		name = 'network_normal'
	else:
		name = 'network' + '_' + name

	with tf.variable_scope(name):
		l_layers = []
		l_layers.append(data)
		l_layers.append(Layers.conv_2d(l_layers[-1], 64, [3, 3], tfv_train_phase = tfv_train_phase, act_first = False, name_scope = 'conv_1'))
		#l_layers.append(Layers.maxpool_2d(l_layers[-1], [3, 3], [2, 2], name_scope = 'maxpool_1'))

		split_depth = [2,3,3,3]
		for i in range(split_depth[0]):
			if i == 0:
				l_layers.append(Layers.resnekt_bottleneck(l_layers[-1], 48, 96, 128, tfv_train_phase = tfv_train_phase, downsample_first = False, name_scope = 'bottleneck_1_%d' % (i + 1)))
			else:
				l_layers.append(Layers.resnekt_bottleneck(l_layers[-1], 48, 96, 128, tfv_train_phase = tfv_train_phase, name_scope = 'bottleneck_1_%d' % (i + 1)))
		for i in range(split_depth[1]):
			l_layers.append(Layers.resnekt_bottleneck(l_layers[-1], 96, 192, 256, tfv_train_phase = tfv_train_phase, name_scope = 'bottleneck_2_%d' % (i + 1)))
		for i in range(split_depth[2]):
			l_layers.append(Layers.resnekt_bottleneck(l_layers[-1], 128, 256, 384, tfv_train_phase = tfv_train_phase, name_scope = 'bottleneck_3_%d' % (i + 1)))
		for i in range(split_depth[3]):
			if i != split_depth[3] - 1:
				l_layers.append(Layers.resnekt_bottleneck(l_layers[-1], 192, 384, 512, tfv_train_phase = tfv_train_phase, name_scope = 'bottleneck_4_%d' % (i + 1)))
			else:
				l_layers.append(Layers.resnekt_bottleneck(l_layers[-1], 192, 384, 512, tfv_train_phase = tfv_train_phase, act_last = True, name_scope = 'bottleneck_4_%d' % (i + 1)))

		l_layers.append(Layers.avgpool_2d(l_layers[-1], [4, 4], [1, 1], padding = 'VALID', is_gap = True, name_scope = 'avgpool'))
		l_layers.append(Layers.fc(l_layers[-1], NUM_CLASSES, act_last = False, name_scope = 'fc_out'))
		outputs = tf.identity(l_layers[-1], 'final_dense')

	return outputs



def softmax(x):
	x_row_max = x.max(axis=-1)
	x_row_max = x_row_max.reshape(list(x.shape)[:-1]+[1])
	x = x - x_row_max
	x_exp = np.exp(x)
	x_exp_row_sum = x_exp.sum(axis=-1).reshape(list(x.shape)[:-1]+[1])
	sm = x_exp / x_exp_row_sum
	return sm


def initialize():
	global tfob_sess, feed, outputs
	b_gpu_enabled = False
	l_devices = device_lib.list_local_devices()
	for i in range(len(l_devices)):
		if l_devices[i].device_type == 'GPU':
			if l_devices[i].memory_limit > 2 * 1024 * 1024 * 1024 :
				b_gpu_enabled = True
				break

	str_last_ckpt = tf.train.latest_checkpoint(flag_log_dir)
	if str_last_ckpt is None:
		return

	with tf.Graph().as_default(), tf.device('/cpu:0'):
		tfv_train_phase = tf.Variable(True, trainable = False, name = 'var_train_phase', dtype = tf.bool, collections = [])

		feed, data = construct_data()
		outputs = network(data, tfv_train_phase, 'cifar_conv')

		tfob_variable_averages = tf.train.ExponentialMovingAverage(0.9, name = 'avg_variable')
		tfop_variable_averages_apply = tfob_variable_averages.apply(tf.trainable_variables())
		tfob_saver_ema = tf.train.Saver(tfob_variable_averages.variables_to_restore())
		if b_gpu_enabled == True:
			tfob_sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, gpu_options = tf.GPUOptions(allow_growth = True, per_process_gpu_memory_fraction = 0.95)))
		else:
			tfob_sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, device_count = {'GPU': 0}))
		tfob_sess.run(tf.global_variables_initializer())

		if str_last_ckpt is not None:
			tfob_sess.run(tfv_train_phase.assign(False))
			tfob_saver_ema.restore(tfob_sess, str_last_ckpt)
		else:
			return


def run(image):
	global tfob_sess, feed, outputs
	image_data = np.reshape(np.array(image, dtype = np.uint8), (32,32))
	dict_feeder = {feed : image_data}
	result = tfob_sess.run(outputs, dict_feeder)
	l_score_conf = np.squeeze(softmax(result)).tolist()
	top1 = max(l_score_conf)*100
	category = l_score_conf.index(max(l_score_conf)) + 1
	print('The category is %d, the confidence is %.2f%%.' % (category, top1))
	rsltList = [category, top1]
	return rsltList


def szTest(List):
    print (List)
    
    IntegerList = np.random.randint(5, size=(4)).tolist()
    return IntegerList


if __name__ == '__main__':
	initialize()
	for i in range(10):
		image = np.reshape(cv2.imread('05_00099_3.jpg', cv2.IMREAD_GRAYSCALE), (-1)).tolist()
		run(image)
	szTest([])
